package com.lw.leetcode.sort.b;

import java.util.PriorityQueue;

/**
 * Created with IntelliJ IDEA.
 * <p>
 * 378. 有序矩阵中第 K 小的元素
 *
 * @author liw
 * @version 1.0
 * @date 2021/8/17 14:09
 */
public class KthSmallest {

    public static void main(String[] args) {
        KthSmallest test = new KthSmallest();

        //  13
//        int[][] arr = {{1, 5, 9}, {10, 11, 13}, {12, 13, 15}};
//        int k = 8;

        //  2
//        int[][] arr = {{1, 2}, {1, 3}};
//        int k = 3;

        // -5
        int[][] arr = {{-5}};
        int k = 1;

        int count = test.kthSmallest(arr, k);
        System.out.println(count);
    }



    public int kthSmallest(int[][] matrix, int k) {
        int length = matrix.length;
        PriorityQueue<Node> queue = new PriorityQueue<>(length, (a, b) -> Integer.compare(a.value, b.value));
        queue.add(new Node(matrix[0][0], 0, 0));
        int[] arr = new int[length * length];
        arr[0] = 1;
        while (k > 1) {
            Node node = queue.poll();
            int x = node.x;
            int y = node.y;
            if (y + 1 < length && arr[x * length + y + 1] == 0) {
                arr[x * length + y + 1] = 1;
                queue.add(new Node(matrix[x][y + 1], x, y + 1));
            }
            if (x + 1 < length && arr[(x + 1) * length + y] == 0) {
                arr[(x + 1) * length + y] = 1;
                queue.add(new Node(matrix[x + 1][y], x + 1, y));
            }
            k--;
        }
        return queue.poll().value;
    }

    class Node {
        private int value;
        private int x;
        private int y;

        public Node(int value, int x, int y) {
            this.value = value;
            this.x = x;
            this.y = y;
        }
    }

}
